#include "solution.h"
#include <iostream>
#include <algorithm>
#include <vector>
#include <unordered_map>
#include <fstream>
#include <string>
#include <chrono> // 包含时间库
#include <thread>
#include <functional>
#include <set>
#include <cmath>
#include <mutex>
#include <queue>

Solution::Solution(size_t n, size_t m, const std::vector<std::vector<size_t>> &pre_ids, const std::vector<long long> &cal_use_times, const std::vector<long long> &trans_use_times)
    : dag(m, pre_ids), n(n), m(m), op_result_card(m), last_trans_end_t(m), prio(m), cal_cost(cal_use_times), trans_cost(trans_use_times),
      unordered_q(m), q_free_t(n), accessed(m)
{
    best_cost = -1;
    this->n = n;
    this->m = m;
    for (int i = 0; i < m; i++)
    {
        cal_cost[i] = cal_use_times[i];
        trans_cost[i] = trans_use_times[i];
        prio[i] = 0;
    }

    st = 0;
    en = 0;
    // 计算优先级
    for (int i = 0; i < m; i++)
        if (dag.out_degree[i] == 0)
        {
            unordered_q[en++] = i;
        }
    while (st < en)
    {
        size_t op = unordered_q[st++];
        for (int i = dag.rheads[op]; i != -1; i = dag.rnex[i])
        {
            size_t pre_op = dag.rto[i];
            prio[pre_op] = std::max(prio[op] + cal_cost[op], prio[pre_op]);
            dag.out_degree[pre_op]--;
            if (dag.out_degree[pre_op] == 0)
            {
                unordered_q[en++] = pre_op;
            }
        }
    }
}

long long Solution::theoretical_best()
{
    for (int i = 0; i < m; i++)
        dag.in_degree[i] = 0;
    for (int i = 0; i < m; i++)
    {
        for (int j = dag.heads[i]; j != -1; j = dag.nex[j])
        {
            dag.in_degree[dag.to[j]]++;
        }
    }
    st = 0;
    en = 0;
    for (int i = 0; i < m; i++)
        if (dag.in_degree[i] == 0)
        {
            unordered_q[en++] = i;
        }
    long long cost = 0;
    std::vector<long long> op_exec_time(m, 0);
    while (st < en)
    {
        size_t op = unordered_q[st++];
        op_exec_time[op] += cal_cost[op];
        cost = std::max(cost, op_exec_time[op]);
        long long start_t = 0;
        for (int i = dag.heads[op]; i != -1; i = dag.nex[i])
        {
            size_t pre_op = dag.to[i];
            dag.in_degree[pre_op]--;
            op_exec_time[pre_op] = std::max(op_exec_time[pre_op], start_t);
            if (dag.in_degree[pre_op] == 0)
            {
                unordered_q[en++] = pre_op;
            }
        }
    }
    return cost;
}

size_t Solution::get_random_operator()
{
    std::swap(unordered_q[st], unordered_q[st + rand() % (en - st)]);
    q.erase({-prio[unordered_q[st]], unordered_q[st]});
    return unordered_q[st++];
}

size_t Solution::get_prio_operator()
{
    if (is_first_round)
    {
        auto it = q.begin();
        size_t op = it->second;
        q.erase(it);
        return op;
    }
    auto it = q.begin();
    std::vector<size_t> ops;
    for (int i = 0; i < 5 && it != q.end(); i++, it++)
    {
        ops.push_back(it->second);
    }
    size_t op = ops[rand() % ops.size()];
    q.erase({-prio[op], op});

    return op;
}

size_t Solution::get_nex_operator()
{
    return get_prio_operator();
}

size_t Solution::get_random_card()
{
    return rand() % n;
}

size_t Solution::get_cal_t_min_card(size_t op)
{
    size_t card;
    long long cal_end_min = -1;
    std::vector<size_t> v;
    for (int i = dag.rheads[op]; i != -1; i = dag.rnex[i])
    {
        size_t pre_op = dag.rto[i];
        v.emplace_back(pre_op);
    }
    sort(v.begin(), v.end(), [&](int i, int j)
         { return last_trans_end_t[i] < last_trans_end_t[j]; });
    int same = 0;
    for (int i = 0; i < n; i++)
    {
        long long tmp = q_free_t[i];
        for (auto pre_op : v)
        {
            if (last_trans_end_t[pre_op] > tmp)
            {
                tmp = last_trans_end_t[pre_op];
            }
            size_t result_card = op_result_card[pre_op];
            tmp += ((result_card == i) ? 0 : trans_cost[pre_op]);
        }
        tmp += cal_cost[op];
        if (cal_end_min == -1 || tmp < cal_end_min)
        {
            same = 1;
            cal_end_min = tmp;
            card = i;
        }
        else if (tmp == cal_end_min && (!is_first_round))
        {
            same++;
            if (rand() % same == 0)
                card = i;
        }
    }
    return card;
}

void Solution::work()
{
    for (int i = 0; i < n; i++)
    {
        q_free_t[i] = 0;
    }
    for (int i = 0; i < m; i++)
    {
        dag.in_degree[i] = 0;
        accessed[i] = false;
        last_trans_end_t[i] = 0;
    }
    for (int i = 0; i < m; i++)
    {
        for (int j = dag.heads[i]; j != -1; j = dag.nex[j])
        {
            dag.in_degree[dag.to[j]]++;
        }
    }
    st = 0;
    en = 0;
    q.clear();
    for (int i = 0; i < m; i++)
        if (dag.in_degree[i] == 0)
        {
            unordered_q[en++] = i;
            q.insert({-prio[i], i});
        }
    cost = 0;
    res.clear();
    while (!q.empty())
    {
        size_t op = get_nex_operator();

        if (accessed[op])
            continue;
        accessed[op] = true;
        size_t card = get_user_free_card(op);
        long long start_t = q_free_t[card];
        std::vector<size_t> v;
        for (int i = dag.rheads[op]; i != -1; i = dag.rnex[i])
        {
            size_t pre_op = dag.rto[i];
            v.emplace_back(pre_op);
        }
        sort(v.begin(), v.end(), [&](int i, int j)
             { return last_trans_end_t[i] < last_trans_end_t[j]; });
        for (auto pre_op : v)
        {
            if (last_trans_end_t[pre_op] > start_t)
            {
                start_t = last_trans_end_t[pre_op];
            }
            size_t result_card = op_result_card[pre_op];
            start_t += ((result_card == card) ? 0 : trans_cost[pre_op]);
        }
        op_result_card[op] = card;
        last_trans_end_t[op] = start_t + cal_cost[op];
        cost = std::max(cost, last_trans_end_t[op]);
        q_free_t[card] = last_trans_end_t[op];
        for (int i = dag.rheads[op]; i != -1; i = dag.rnex[i])
        {

            size_t pre_op = dag.rto[i];
            last_trans_end_t[pre_op] = start_t;
            op_result_card[pre_op] = card;
        }
        res.push_back({op, card});
        for (int i = dag.heads[op]; i != -1; i = dag.nex[i])
        {
            size_t nex_op = dag.to[i];
            dag.in_degree[nex_op]--;
            if (dag.in_degree[nex_op] == 0)
            {
                unordered_q[en++] = nex_op;
                q.insert({-prio[nex_op], nex_op});
            }
        }
    }
    if (best_cost == -1 || cost < best_cost)
    {
        cnt_from_last_optimize = 0;
        best_cost = cost;
        best_res = res;
    }
    else
        cnt_from_last_optimize++;
    is_first_round = false;
}

void Solution::show_best_res(const std::vector<size_t> &ids)
{
    for (int i = 0; i < best_res.size(); i++)
    {
        std::cout << "op_id:" << best_res[i].first << " card_id:" << ids[best_res[i].second] << std::endl;
    }
    std::cout << "best_cost:" << best_cost << std::endl;
}

std::pair<long long, std::vector<std::pair<size_t, size_t>>> Solution::get_best_res()
{
    return {best_cost, best_res};
}

void Solution::merge_best_res(Solution &sol1, Solution &sol2)
{
    if (sol1.best_cost == -1)
    {
        sol1.best_cost = sol2.best_cost;
        sol1.best_res = sol2.best_res;
    }
    else if (sol2.best_cost != -1 && sol2.best_cost < sol1.best_cost)
    {
        sol1.best_cost = sol2.best_cost;
        sol1.best_res = sol2.best_res;
    }
}
void Solution::optimize_res(bool optimize_best_res)
{
    // TODO 局部优化
    if (optimize_best_res)
        res = this->best_res;
    else
        res = this->res;
    cost = 0;
    for (int i = 0; i < n; i++)
    {
        q_free_t[i] = 0;
    }
    bool exchanged = false;
    if (rand() % 2 == 0)
    {
        int index = rand() % (res.size() - 1);
        size_t op1 = res[index].first;
        size_t op2 = res[index + 1].first;
        bool is_pre = false;
        for (int j = dag.heads[op1]; j != -1; j = dag.nex[j])
        {
            if (dag.to[j] == op2)
            {
                is_pre = true;
                break;
            }
        }
        if (!is_pre)
        {
            std::swap(res[index], res[index + 1]);
            exchanged = true;
        }
    }
    if (!exchanged)
    {
        int index = rand() % (res.size());
        size_t new_card = rand() % (n - 1);
        if (new_card >= res[index].second)
            new_card++;
        res[index].second = new_card;
    }
    for (int i = 0; i < n; i++)
    {
        q_free_t[i] = 0;
    }
    for (int i = 0; i < res.size(); i++)
    {
        size_t op = res[i].first;
        size_t card = res[i].second;
        long long start_t = q_free_t[card];
        std::vector<size_t> v;
        for (int i = dag.rheads[op]; i != -1; i = dag.rnex[i])
        {
            size_t pre_op = dag.rto[i];
            v.emplace_back(pre_op);
        }
        sort(v.begin(), v.end(), [&](int i, int j)
             { return last_trans_end_t[i] < last_trans_end_t[j]; });
        for (auto pre_op : v)
        {
            if (last_trans_end_t[pre_op] > start_t)
            {
                start_t = last_trans_end_t[pre_op];
            }
            size_t result_card = op_result_card[pre_op];
            start_t += ((result_card == card) ? 0 : trans_cost[pre_op]);
        }
        op_result_card[op] = card;
        last_trans_end_t[op] = start_t + cal_cost[op];
        cost = std::max(cost, last_trans_end_t[op]);
        q_free_t[card] = last_trans_end_t[op];
        for (int i = dag.rheads[op]; i != -1; i = dag.rnex[i])
        {
            size_t pre_op = dag.rto[i];
            last_trans_end_t[pre_op] = start_t;
            op_result_card[pre_op] = card;
        }
    }
    if (optimize_best_res)
    {
        if (cost < best_cost)
        {
            this->best_cost = cost;
            this->best_res = res;
        }
    }
    else
    {
        if (cost < this->cost)
        {
            this->cost = cost;
            this->res = res;
        }
    }
}
std::mutex file_mutex;

void write_res_to_file(const Solution &sol, std::string output_path)
{
    std::lock_guard<std::mutex> lock(file_mutex);
    auto res = sol.get_res();
    auto cost = sol.get_cost();
    // 追加写
    sort(res.begin(), res.end(), [](std::pair<size_t, size_t> &a, std::pair<size_t, size_t> &b)
         { return a.second < b.second; });
    std::ofstream fout(output_path, std::ios::app);
    fout << cost << std::endl;
    fout << '[';
    for (int i = 0; i < res.size(); i++)
    {
        fout << res[i].first;
        if (i != res.size() - 1)
            fout << ',';
    }
    fout << ']' << std::endl;
}

// 使用 std::thread 替代 pthread（更现代、更安全）
static void run_in_threads(const std::vector<std::function<void()>> &tasks)
{
    std::vector<std::thread> threads;
    threads.reserve(tasks.size());
    for (auto &task : tasks)
    {
        threads.emplace_back(task);
    }
    for (auto &t : threads)
    {
        if (t.joinable())
        {
            t.join();
        }
    }
}

void MergeSolutions(std::vector<Solution> &solutions)
{
    // 合并最优解到第0个
    for (size_t i = 1; i < solutions.size(); ++i)
    {
        Solution::merge_best_res(solutions[0], solutions[i]);
    }

    // 重置其他解，继承当前最优
    for (size_t i = 1; i < solutions.size(); ++i)
    {
        solutions[i].cp_best_res(solutions[0]);
    }
}

void GernerateInitialSolution(std::vector<Solution> &solutions, double time_limit_seconds)
{
    auto start = std::chrono::high_resolution_clock::now();

    do
    {
        std::vector<std::function<void()>> tasks;
        for (auto &sol : solutions)
        {
            tasks.emplace_back([&sol]()
                               { sol.work(); });
        }
        run_in_threads(tasks);

        auto now = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double> elapsed = now - start;

        if (elapsed.count() >= time_limit_seconds)
            break;

    } while (true);
    MergeSolutions(solutions);
}

static void OptimizeBestSolution(
    std::vector<Solution> &solutions,
    double time_limit_seconds)
{
    auto start = std::chrono::high_resolution_clock::now();

    do
    {
        std::vector<std::function<void()>> tasks;
        for (auto &sol : solutions)
        {
            tasks.emplace_back([&sol]()
                               { sol.optimize_res(); });
        }
        run_in_threads(tasks);

        auto now = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double> elapsed = now - start;
        if (elapsed.count() >= time_limit_seconds)
            break;

    } while (true);
    MergeSolutions(solutions);
}

std::vector<std::pair<size_t, size_t>> solve(
    size_t n,
    size_t m,
    const std::vector<long long> &cal_use_times,
    const std::vector<long long> &trans_use_times,
    const std::vector<size_t> &ids,
    const std::vector<std::vector<size_t>> &graph,
    int thread_size)
{
    double run_upper_time = 60 * ((double)m / (double)50000);
    double inital_solution_t = run_upper_time * 0.8;
    double optimize_solution_t = run_upper_time - inital_solution_t;
    auto start = std::chrono::high_resolution_clock::now();
    std::vector<Solution> solutions;
    if (n == 1 or m == 1)
    {
        Solution sol(n, m, graph, cal_use_times, trans_use_times);
        sol.work();
        return sol.get_best_res().second;
    }
    for (int i = 0; i < thread_size; i++)
    {
        solutions.push_back(Solution(n, m, graph, cal_use_times, trans_use_times));
    }

    // 一阶段
    GernerateInitialSolution(solutions, inital_solution_t);
    // 二阶段
    OptimizeBestSolution(solutions, optimize_solution_t);

    return solutions[0].get_best_res().second;
}

void InitAllNodes(const std::vector<Node *> &all_nodes, size_t card_num, std::vector<long long> &cal_use_times,
                  std::vector<long long> &trans_use_times, std::vector<size_t> &ids, std::vector<std::vector<size_t>> &graph)
{
    std::unordered_map<size_t, size_t> id_to_index;
    for (int i = 0; i < all_nodes.size(); i++)
    {
        ids.emplace_back(all_nodes[i]->id());
        id_to_index[all_nodes[i]->id()] = i;
    }
    size_t n = card_num;
    size_t m = all_nodes.size();
    for (auto &node : all_nodes)
    {
        std::vector<size_t> unique_pre_ids;
        cal_use_times.emplace_back(node->exec_time());
        trans_use_times.emplace_back(node->transfer_time());
        for (auto &pre : node->inputs())
        {
            unique_pre_ids.emplace_back(pre->id());
        }
        sort(unique_pre_ids.begin(), unique_pre_ids.end());
        int len = 0;
        for (int j = 0; j < unique_pre_ids.size(); j++)
        {
            if (j == 0 || unique_pre_ids[j] != unique_pre_ids[j - 1])
            {
                unique_pre_ids[len++] = unique_pre_ids[j];
            }
        }
        while (unique_pre_ids.size() > len)
            unique_pre_ids.pop_back();
        for (int j = 0; j < unique_pre_ids.size(); j++)
            unique_pre_ids[j] = id_to_index[unique_pre_ids[j]];
        graph[id_to_index[node->id()]] = unique_pre_ids;
    }
}

std::vector<std::pair<size_t, size_t>>
ExecuteOrder(const std::vector<Node *> &all_nodes, size_t card_num)
{
    std::vector<long long> cal_use_times;
    std::vector<long long> trans_use_times;
    std::vector<size_t> ids;
    std::vector<std::vector<size_t>> graph(all_nodes.size());
    InitAllNodes(all_nodes, card_num, cal_use_times, trans_use_times, ids, graph);
    return solve(card_num, all_nodes.size(), cal_use_times, trans_use_times, ids, graph, 4);
}